package com.xuanwu.ai.flow.utils;

import com.alibaba.fastjson.JSONObject;
import com.xuanwu.ai.core.domain.SessionInfo;
import com.xuanwu.ai.core.exception.DataAccessException;
import com.xuanwu.ai.flow.constant.GlobalConstant;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jws;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.impl.DefaultClaims;
import io.jsonwebtoken.impl.TextCodec;

import java.util.Date;
import java.util.Map;
import java.util.Optional;

public class JwtUtil {

    /**
     * plattoken签名密钥
     */
    private final static byte[] PLAT_SECURITY_KEY = {0x36, 0x31, 0x33, 0x30, 0x32, 0x32, 0x32, 0x32, 0x32};

    /**
     * ai-flow签名密钥
     */
    private final static byte[] FLOW_SECURITY_KEY = {0xF, 0xC, 0x33, 0x30, 0x32, 0x32, 0x32, 0xA, 0xB};

    public static String wrap(SessionInfo sessionInfo, Date expires) {
        Claims claims = new DefaultClaims();
        claims.setExpiration(expires);
        claims.put("LoginUser", sessionInfo);
        return Jwts.builder()
                .setClaims(claims)
                .signWith(SignatureAlgorithm.HS256, TextCodec.BASE64.encode(FLOW_SECURITY_KEY))
                .compact();
    }

    /**
     * 解析token
     *
     * @param token
     * @return
     * @throws Exception
     */
    public static SessionInfo unwrap(String token, String flag) {
        try {
            return unwrap(token, 0, flag);
        } catch (Exception ex) {
            throw new DataAccessException("jwt token无效", ex);
        }
    }

    public static SessionInfo unwrap(String token, int allowedSkewSeconds, String flag) {
        byte[] securityKey = GlobalConstant.PLATFORM_UNWRAP_FLAG.equals(flag) ? PLAT_SECURITY_KEY : FLOW_SECURITY_KEY;
        Jws<Claims> jws = Jwts.parser()
                .setSigningKey(TextCodec.BASE64.encode(securityKey))
                .setAllowedClockSkewSeconds(allowedSkewSeconds)
                .parseClaimsJws(token);

        Claims claims = jws.getBody();

        Map map = claims.get("LoginUser", Map.class);
        JSONObject jsonObject = new JSONObject(map);
        SessionInfo sessionInfo = jsonObject.toJavaObject(SessionInfo.class);

        if (sessionInfo != null) {
            return sessionInfo;
        }
        throw new RuntimeException("token格式有误，找不到登录载体:[" + Optional.ofNullable(token).orElse("") + "]");
    }
}
